package com.cspm.common.proxy;

import com.alibaba.druid.pool.DruidDataSource;
import com.cspm.gen.enums.DriverCLassEnum;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 动态数据源实现读写分离
 *
 * @author CHENQUAN
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    // 写数据源
    private Object writeDataSource;
    // 多个读数据源
    private List<Object> readDataSources;
    // 读数据源个数
    private int readDataSourceSize;
    // 获取读数据源方式，0：随机，1：轮询
    private int readDataSourcePollPattern = 0;

    private final AtomicLong counter = new AtomicLong(0);

    private static final Long MAX_POOL = Long.MAX_VALUE;

    private final Lock lock = new ReentrantLock();
    //自定义数据源
    private static javax.sql.DataSource customerDataSource = null;

    @Override
    public void afterPropertiesSet() {
        if (customerDataSource != null) {
            setDefaultTargetDataSource(customerDataSource);
        } else {
            if (this.writeDataSource == null) {
                throw new IllegalArgumentException("Property 'writeDataSource' is required");
            }
            setDefaultTargetDataSource(writeDataSource);
            Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
            targetDataSources.put(DynamicDataSourceGlobal.WRITE.name(), writeDataSource);
            if (this.readDataSources == null) {
                readDataSourceSize = 0;
            } else {
                for (int i = 0; i < readDataSources.size(); i++) {
                    targetDataSources.put(DynamicDataSourceGlobal.READ.name() + i, readDataSources.get(i));
                }
                readDataSourceSize = readDataSources.size();
            }
            setTargetDataSources(targetDataSources);
        }
        super.afterPropertiesSet();
    }

    /**
     * 确定当前查找的数据源
     */
    @Override
    protected Object determineCurrentLookupKey() {
        DynamicDataSourceGlobal dynamicDataSourceGlobal = DynamicDataSourceHolder.getDataSource();
        if (dynamicDataSourceGlobal == null
                || dynamicDataSourceGlobal == DynamicDataSourceGlobal.WRITE
                || readDataSourceSize <= 0) {
            //默认为写数据源
            return DynamicDataSourceGlobal.WRITE.name();
        }

        int index = 1;
        if (readDataSourcePollPattern == 1) {
            // 轮询方式
            long currValue = counter.incrementAndGet();
            if ((currValue + 1) >= MAX_POOL) {
                try {
                    lock.lock();
                    if ((currValue + 1) >= MAX_POOL) {
                        counter.set(0);
                    }
                } finally {
                    lock.unlock();
                }
            }
            index = (int) (currValue % readDataSourceSize);
        } else {
            // 随机方式
            index = ThreadLocalRandom.current().nextInt(0, readDataSourceSize);
        }
        return dynamicDataSourceGlobal.name() + index;
    }

    public void setWriteDataSource(Object writeDataSource) {
        this.writeDataSource = writeDataSource;
    }

    public void setReadDataSources(List<Object> readDataSources) {
        this.readDataSources = readDataSources;
    }

    public void setReadDataSourcePollPattern(int readDataSourcePollPattern) {
        this.readDataSourcePollPattern = readDataSourcePollPattern;
    }

    public static javax.sql.DataSource getCustomerDataSource() {
        return customerDataSource;
    }

    public static void setCustomerDataSource(DriverCLassEnum driverCLassEnum, String server, Integer port, String database, String username, String password) {
        DruidDataSource druidDataSource = new DruidDataSource();
        druidDataSource.setUrl(driverCLassEnum.getUrl(server, port, database));
        druidDataSource.setUsername(username);
        druidDataSource.setPassword(password);
        customerDataSource = druidDataSource;
    }
}